import collections
import copy
import datetime
import logging
import pytz

from absl import app, flags

import coin.proto.coin_order_enums_pb2 as order_enums_pb2
import coin.support.proto_log.logic.util as proto_log_util
from coin.base.datetime_util import (
    convert_string_to_datetime,
    get_dt_ranges,
    to_datetime,
    to_timestamp_int)
from coin.base.logging import init_logging
from coin.base.param_util import split_string
from coin.exchange.base.strategy_logger import (
    enumerate_logs as enumerate_strat_logs)
from coin.flow.archive import ArchiveReader
from coin.proto.coin_market_query_pb2 import (
    ExchangeTicker,
    ProductTicker)
from coin.proto.coin_order_gateway_pb2 import (
    OrderGatewayLog,
    OrderEvent)
from coin.proto.coin_request_pb2 import StrategyRequestProto
from coin.proto.coin_strategy_pb2 import StrategyLog
from coin.support.pta.util.info_util import get_strat_info_list
from xunkemgmt_client.client.util.query_util import query_strategies


FLAGS = flags.FLAGS


WorkingOrderInfo = collections.namedtuple('WorkingOrderInfo', [
    'timestamp',  # int
    'market_type',  # string
    'exchange',  # string
    'symbol',  # string
    'working_order',  # [OrderGatewayLog]
    'ticker',  # ExchangeTicker
])

OrderKey = collections.namedtuple('OrderKey', [
    'market_type',
    'exchange',
    'owner',
    'symbol',
    'proc_order_id',
    'machine',
    'pid',
])

TickerKey = collections.namedtuple('TickerKey', [
    'market_type', 'exchange', 'symbol'
])


def _gen_interval_iter(start_dt, end_dt, interval):
  assert isinstance(start_dt, datetime.datetime)
  assert isinstance(end_dt, datetime.datetime)
  assert isinstance(interval, datetime.timedelta)
  next_dt = start_dt + interval
  while next_dt < end_dt:
    yield next_dt
    next_dt += interval
  yield end_dt


def _validate_og_log(strat_log):
  assert strat_log.type == StrategyLog.OG_LOG
  assert strat_log.HasField('process')
  process = strat_log.process
  assert process.HasField('machine')
  assert process.HasField('pid')
  og_log = strat_log.og_log
  assert og_log.type == OrderGatewayLog.ORDER_EVENT
  assert og_log.HasField('account_request')
  account_request = og_log.account_request
  assert account_request.HasField('market_type')
  assert account_request.HasField('exchange')
  assert account_request.HasField('owner')
  order_event = og_log.event
  assert order_event.HasField('symbol')
  assert order_event.HasField('event_time')
  assert order_event.HasField('order_state')
  assert order_event.HasField('proc_order_id')
  if order_event.order_state == order_enums_pb2.PENDING_ORDER:
    assert order_event.HasField('order_created_time')


def _validate_ticker(exchange_ticker):
  assert isinstance(exchange_ticker, ExchangeTicker)
  assert exchange_ticker.HasField('market_type')
  assert exchange_ticker.HasField('exchange')
  assert exchange_ticker.HasField('query_ts')
  assert len(exchange_ticker.each_ticker) == 1
  product_ticker = exchange_ticker.each_ticker[0]
  assert product_ticker.HasField('symbol')


def _gen_order_key(strat_log):
  _validate_og_log(strat_log)
  process = strat_log.process
  account_request = strat_log.og_log.account_request
  order_event = strat_log.og_log.event
  key = OrderKey(
      market_type=account_request.market_type,
      exchange=account_request.exchange,
      owner=account_request.owner,
      symbol=order_event.symbol,
      proc_order_id=order_event.proc_order_id,
      machine=process.machine,
      pid=process.pid)
  return key


def _gen_ticker_key(exchange_ticker):
  _validate_ticker(exchange_ticker)
  symbol = exchange_ticker.each_ticker[0].symbol
  key = TickerKey(
      market_type=exchange_ticker.market_type,
      exchange=exchange_ticker.exchange,
      symbol=symbol)
  return key


def _print_stats(stats):
  for key, value in stats.items():
    print('market_type: %s, exchange: %s, symbol: %s' % (
        key.market_type, key.exchange, key.symbol))
    for order_info in value:
      assert isinstance(order_info, WorkingOrderInfo)
      print('timestamp: %s' % to_datetime(order_info.timestamp))
      if order_info.working_order is not None:
        for each_order in order_info.working_order:
          order_event = each_order.event
          print('order_price: %s, order_qty: %s, order_side: %s, created_time: %s' % (
              order_event.order_price,
              order_event.order_qty,
              order_enums_pb2.OrderSide.Name(order_event.order_side),
              to_datetime(order_event.order_created_time)))
      if order_info.ticker is not None:
        print('price: %s time: %s' % (
            order_info.ticker.each_ticker[0].last,
            to_datetime(order_info.ticker.query_ts)))


class WorkingOrderDumper(object):
  def __init__(self, *,
               start_dt, end_dt, interval, strat_list,
               strat_log_root_dir,
               working_order_limit_sec=3000):
    self._start_dt = start_dt
    self._end_dt = end_dt
    self._strat_list = strat_list
    self._strat_log_root_dir = strat_log_root_dir
    self._last_ts = 0
    self._time_point_iter = _gen_interval_iter(start_dt, end_dt, interval)
    self._next_time_point = next(self._time_point_iter)
    self._working_order_dict = {}
    self._ticker_dict = {}
    self._stats = collections.defaultdict(list)
    self._working_order_limit_sec = working_order_limit_sec

  def replay(self):
    log_files = []
    dt_range = proto_log_util._get_date_range(self._start_dt, self._end_dt)
    start_ts = to_timestamp_int(self._start_dt)
    end_ts = to_timestamp_int(self._end_dt)
    strat_infos = get_strat_info_list(
        start_date=self._start_dt.date(),
        end_date=self._end_dt.date(),
        strategy_list=self._strat_list,
        strat_log_root_dir=self._strat_log_root_dir)
    for dt in dt_range:
      for strat_info in strat_infos:
        log_dir = proto_log_util._gen_log_dir(
            log_root=self._strat_log_root_dir,
            trading_date=dt,
            machine=strat_info.machine)
        strat_req = StrategyRequestProto(strategy_name=strat_info.strategy_name)
        log_files += enumerate_strat_logs(strat_req, dt, log_dir)
    with ArchiveReader.from_file(log_files) as reader:
      while True:
        ret = reader.read()
        if ret is None:
          break
        ts = ret.timestamp
        if ts < start_ts:
          continue
        if ts > end_ts:
          break
        self.on_log(ts, ret.data)

  def get_stats(self):
    for time_point in self._time_point_iter:
      self._add_stats(time_point)
    return copy.deepcopy(self._stats)

  def on_log(self, ts, log):
    if ts < self._last_ts:
      return
    dt = datetime.datetime.utcfromtimestamp(ts / 10**9).replace(tzinfo=pytz.UTC)
    while dt > self._next_time_point:
      self._add_stats(self._next_time_point)
      try:
        self._next_time_point = next(self._time_point_iter)
      except StopIteration:
        return

    pb = StrategyLog()
    pb.ParseFromString(log)
    if pb.type == StrategyLog.OG_LOG:
      self._handle_og_log(ts, pb)
    elif pb.type == StrategyLog.TICKER:
      self._handle_ticker(ts, pb)

  def _update_ticker_dict(self, ticker):
    assert isinstance(ticker, ExchangeTicker)
    ticker_key = _gen_ticker_key(ticker)
    old_ticker = self._ticker_dict.get(ticker_key)
    if old_ticker is None or \
       ticker.query_ts > old_ticker.query_ts:
      self._ticker_dict[ticker_key] = ticker

  def _handle_og_log(self, ts, strat_log):
    if strat_log.og_log.type != OrderGatewayLog.ORDER_EVENT:
      return
    order_key = _gen_order_key(strat_log)
    order_event = strat_log.og_log.event
    if order_event.order_state == order_enums_pb2.PENDING_ORDER:
      self._working_order_dict[order_key] = strat_log.og_log
    elif order_event.order_state == order_enums_pb2.DEAD_ORDER:
      self._working_order_dict.pop(order_key, None)
    if order_event.type == OrderEvent.ORDER_FILLED:
      assert order_event.HasField('fill_price')
      account_request = strat_log.og_log.account_request
      exchange_ticker = ExchangeTicker(
          market_type=account_request.market_type,
          exchange=account_request.exchange,
          query_ts=order_event.event_time,
          each_ticker=[ProductTicker(
              symbol=order_event.symbol,
              last=order_event.fill_price)])
      self._update_ticker_dict(exchange_ticker)

  def _handle_ticker(self, ts, strat_log):
    assert isinstance(strat_log, StrategyLog)
    assert strat_log.type == StrategyLog.TICKER
    for ticker in strat_log.ticker.each_ticker:
      exchange_ticker = ExchangeTicker(
          market_type=strat_log.ticker.market_type,
          exchange=strat_log.ticker.exchange,
          query_ts=ts,
          each_ticker=[ticker])
      self._update_ticker_dict(exchange_ticker)

  def _gen_working_order_dict(self, dt):
    working_order_dict = collections.defaultdict(list)
    remove_keys = []
    for key in self._working_order_dict.keys():
      working_order = self._working_order_dict[key]
      if to_timestamp_int(dt) - working_order.event.order_created_time > self._working_order_limit_sec * 10**9:
        remove_keys.append(key)
      else:
        prod_key = TickerKey(
            market_type=key.market_type, exchange=key.exchange, symbol=key.symbol)
        working_order_dict[prod_key].append(working_order)
    for key in remove_keys:
      value = self._working_order_dict.pop(key)
      logging.debug('remove outdated order: %s.%s.%s %s\n%s' % (
          key.market_type, key.exchange, key.symbol, dt, value))
    return working_order_dict

  def _gen_ticker_dict(self, dt):
    ticker_dict = {}
    remove_keys = []
    for key in self._ticker_dict.keys():
      ticker = self._ticker_dict[key]
      if to_timestamp_int(dt) - ticker.query_ts > 30 * 10**9:
        remove_keys.append(key)
      else:
        ticker_dict[key] = ticker
    for key in remove_keys:
      value = self._ticker_dict.pop(key)
      logging.debug('remove outdated ticker: %s.%s.%s %s\n%s' % (
          key.market_type, key.exchange, key.symbol, dt, value))
    return ticker_dict

  def _add_stats(self, dt):
    working_order_dict = self._gen_working_order_dict(dt)
    ticker_dict = self._gen_ticker_dict(dt)
    keys = set(list(working_order_dict.keys()) + list(ticker_dict.keys()))
    for key in keys:
      ticker = ticker_dict.get(key)
      working_order = working_order_dict.get(key)
      working_order_info = WorkingOrderInfo(
          timestamp=int(dt.timestamp() * 10**9),
          market_type=key.market_type,
          exchange=key.exchange,
          symbol=key.symbol,
          working_order=working_order,
          ticker=ticker)
      self._stats[key].append(working_order_info)


def main(_):
  init_logging()
  start_time = FLAGS.start_time
  assert start_time, '--start_time must be specified.'
  end_time = FLAGS.end_time
  assert end_time, '--end_time must be specified.'

  start_dt = convert_string_to_datetime(start_time)
  end_dt = convert_string_to_datetime(end_time)
  dt_ranges = get_dt_ranges(start_dt, end_dt, FLAGS.aggregate)

  strategy_universe = query_strategies(business_units=split_string(FLAGS.business_unit),
                                       strategy_groups=split_string(FLAGS.strategy_group),
                                       strategy_names=split_string(FLAGS.strategy_name),
                                       as_proto=True)
  strat_list = list(set([elem.strategy_name for elem in strategy_universe]))
  interval = datetime.timedelta(seconds=FLAGS.interval_sec)

  for start_dt, end_dt in dt_ranges:
    start_td_str = start_dt.strftime('%Y%m%dT%H%M%S')
    end_td_str = end_dt.strftime('%Y%m%dT%H%M%S')
    logging.info('\nRunning for %s-%s ...' % (start_td_str, end_td_str))
    dumper = WorkingOrderDumper(
        start_dt=start_dt,
        end_dt=end_dt,
        interval=interval,
        strat_list=strat_list,
        strat_log_root_dir=FLAGS.root_dir)
    dumper.replay()
    stats = dumper.get_stats()
    #_print_stats(stats)


if __name__ == '__main__':
  flags.DEFINE_string('start_time', None, 'Start time in form of %Y%m%d or %Y%m%dT%H%M%S.')

  flags.DEFINE_string('end_time', None, 'End time in form of %Y%m%d or %Y%m%dT%H%M%S.')

  flags.DEFINE_boolean('aggregate', False, '')

  flags.DEFINE_float('interval_sec', 60, 'interval_sec')

  flags.DEFINE_string('strategy_group', None, '')

  flags.DEFINE_string('strategy_name', None, '')

  flags.DEFINE_string('business_unit', None, 'filter by business_unit')

  flags.DEFINE_string('root_dir', None, '')

  app.run(main)
